import numpy as np
import numpy.linalg as lg
import matplotlib.pyplot as mp
from matplotlib.pylab import *

#---- Fonction definissant la methode de resolution d euler -------- #
def step_euler(y,t,h,f):
    n = np.shape(y)[0]
   
    z = np.zeros([n,1])
    z.shape=(n)
   
    for i in range(n):
        z[i] = y[i] + h * f[i](y,t)
   
    return z


#---- Fonction definissant la methode de resolution du point milieu ---#
def step_point_milieu(y,t,h,f):
    n = np.shape(y)[0]
    z = np.zeros([n,1])
    z.shape=(n)
    p = np.zeros([n,1])
    for i in range(n):
        z[i] = y[i] + (h/2) * f[i](y,t)
    for i in range(n):
        p[i] = f[i](z,t+h/2)
    for i in range(n):
        z[i] = y[i] + h*p[i] 
    return z

#------ Fonction definissant la methode de resolution Heun -----------#
def step_heun(y,t,h,f):
    n = np.shape(y)[0]
    z = np.zeros([n,1])
    z.shape=(n)
    p1 = np.zeros([n,1])
    p1.shape=(n)
    p2 = np.zeros([n,1])
    p2.shape=(n)
    for i in range(n):
        p1[i] = f[i](y,t)
    for i in range(n):
        z[i]  = y[i] + h*p1[i]
    for i in range(n):
        p2[i] = f[i](z,t+h)
    for i in range(n):
        z[i] = y[i] + (h/2) * (p1[i] + p2[i])

    return z


#------ Fonction definissant la methode de resolution de Runge Kutta  --- #
def runge_kutta_step(y,t,h,f):
    n = np.shape(y)[0]
    p1 = np.zeros([n,1])
    p1.shape=(n)
    p2 = np.zeros([n,1])
    p2.shape=(n)
    p3 = np.zeros([n,1])
    p3.shape=(n)
    p4 = np.zeros([n,1])
    p4.shape=(n)
    z = np.zeros([n,1])
    z.shape=(n)
    for i in range(n):
        p1[i] = f[i](y,t)
  
    for i in range(n):
        z[i] = y[i] + h/2 * p1[i]
    
    for i in range(n):
        p2[i] = f[i](z,t+ h/2)
  
    
    for i in range(n):
        z[i] = y[i] + h/2 * p2[i]
    for i in range(n):
        p3[i] = f[i](z,t+ h/2)
    
    for i in range(n):
        z[i] = y[i] + h * p3[i]
    for i in range(n):
        p4[i] = f[i](z,t+ h)
    
    for i in range(n):
        z[i] = y[i] + (1./6.) * h * (p1[i] + 2*p2[i] + 2*p3[i] + p4[i])
        

    return z
                                   
# --- Fonction retournant le vectuer  image du couple (y,t) par la fonction f -----#
def apply_function(f,y,t):
    n = y.shape[0]
    X = np.zeros([n,1])
    X.shape = (n)
    for i in range(n):
        X[i] = f[i](y,t)

    return X


#---- Fonction generique effectuant n pas en appliquant la methode passe en parametre 'meth' a partir de t0 avec un pas de h ----#
def meth_step_n(y0,t0,n,h,f,meth):
    y = y0
    t = t0
    N = np.shape(y)[0]
    #vecteur contenant les y calcule a chaque pas
    func = np.zeros([n+1,N])
    #vecteur contenant les f(y,t) calcule a chaque pas 
    tang = np.zeros([n+1,N])
    
    func[0] = y0
    
    tang[0] = apply_function(f,y0,t0)
    for i in range(n):
        y = meth(y,t,h,f)
        
        func[i+1] = y
        tang[i+1] = apply_function(f,y,t)
        t = t + h
   
    return (func,tang)


#---- Fonction appliquant la methode de resolution passe en parametre 'meth' entre t0 et tf
#----                avec une constante d erreur eps                        -----#
def meth_epsilon(y0,t0,tf,eps,f,meth):
    N = 2
    (X,tang) = meth_step_n(y0,t0,N,(tf-t0)/N,f,meth)
    first = 1   #variable servant a distinguer le premier pas
    old = X[np.shape(X)[0]-1]
    new = X[np.shape(X)[0]-1]
    while( (lg.norm(new - old ) > eps or first == 1) ) :
        first = 0
        
        old = X[np.shape(X)[0]-1]  #enregistrer la derniere valeur du vecteur calcule apres N pas
        N = N*2 # le nombre de pas est double a chaque tour de boucle
        (X,tang) = meth_step_n(y0,t0,N,(tf-t0)/N,f,meth)
        new = X[np.shape(X)[0]-1] 
      
    return (X,tang)

#--- Fonction qui extrait les composantes x d un vecteur ----#
def extract_abcissas(X,n):
    absc = np.zeros([n,1])
    absc.shape = (n)
    for i in range(n):
        absc[i] = X[i][0] 
    return absc
#---- Fonction qui extrait les composantes y d un vecteur ----#
def extract_ordinates(X,n):
    ordi = np.zeros([n,1])
    ordi.shape = (n)
    for i in range(n):
        ordi[i] = X[i][1] 
    return ordi

# Fonction qui trace le champ des tangentes pour une fonction f dans la grille [[xmin,xmax],[ymin,ymax]] ----#
def plot_tang_field(xmin,xmax,ymin,ymax,NX,NY,f):
    
    X = linspace(xmin, xmax, NX)
    Y = linspace(ymin, ymax, NY)
    U = np.ones([NX,NY])
    V = np.ones([NX,NY])
    for i in range(NX):
        for j in range(NY):
            V[i,j] = f([Y[i]],Y[j])
    mp.clf()
    mp.quiver(X,Y,U,V)
    mp.title('Champ de tangentes pour la fonction f(x,y) ')
    mp.xlabel('x ')
    mp.ylabel('y ')
    
    mp.savefig("champ-tangentes.png")                
    mp.show()

#------------------- Test ------------------------------------------#


#------On commence par construire la fonction differentielle definissant le systeme -----------# 

#--- solution exacte y = exp(arctan(t)) -------#
def function_exp_arctan():
    def f(y,t):
        return y[0]/(1+t**2)
   
    return [f]
#------ solution exacte (y1(t),y2(t)) = (cos(t),sin(t) -------#
def function_cos_sin():
    def f(y,t):
        return -y[1]
    def g(y,t):
        return y[0]
    return [f,g]


#------ solution exacte y = K*exp(at)- (b/a) -----#
def function_exp(a,b):
    def f(y,t):
        return a*y+b
    return [f]

# ------- Test des differents methodes et comparaison avec la solution exacte  ------#

#---- Test pour la fonction f = (y/1+t^2) -------#

f = function_exp_arctan()
Y1 = meth_epsilon(np.array([1.0]),0.,100.,0.005,f,step_euler)[0]
Y2 = meth_epsilon(np.array([1.0]),0.,100.,0.05,f,step_heun)[0]
Y3 = meth_epsilon(np.array([1.0]),0.,100.,0.05,f,step_point_milieu)[0]
Y4 = meth_epsilon(np.array([1.0]),0.,100.,0.05,f,runge_kutta_step)[0]

n1 = Y1.shape[0]
n2 = Y2.shape[0]
n3 = Y3.shape[0]
n4 = Y4.shape[0]

X1 = linspace(0.,100.,n1)
X2 = linspace(0.,100.,n2)
X3 = linspace(0.,100.,n3)
X4 = linspace(0.,100.,n4)


S = np.zeros([n4,1])
S.shape = (n4)
for i in range(n4):
    S[i] = math.atan(X4[i])
for i in range(n4):
    S[i] = math.exp(S[i])
mp.clf()

r = mp.plot(X4,S)
b = mp.plot(X1,Y1)
g = mp.plot(X2,Y2)
k = mp.plot(X3,Y3)
m = mp.plot(X4,Y4)



mp.legend((r,b,g,k,m),("Solution exacte","Euler","Heun","Point milieu","Runge Kutta"))
mp.xlabel('axe t')
mp.ylabel('axe y(t)')
mp.title('Resolution de l equation y`=y/(1+t^2)')
mp.savefig("first-equation.png")


#----- Test avec la fonction f = (-y2(t),y1(t)) ------#

mp.clf()
f = function_cos_sin()

Y = meth_epsilon(np.array([1.0,0.0]),0.,30.,0.05,f,runge_kutta_step)[0]
X = linspace(0.,30.,Y.shape[0])

Y1 = extract_abcissas(Y,Y.shape[0])
Y2 = extract_ordinates(Y,Y.shape[0])
a = mp.plot(X,Y1)
b = mp.plot(X,Y2)
mp.title("Resolution de l equation differentielle y` = (-y2(t),y1(t))")
mp.xlabel("axe t")
mp.ylabel("axe y(t)")
mp.legend((a,b),("y1(t) = cos(t)","y2(t) = sin(t)"))
mp.savefig("second-equation.png")


mp.clf()

mp.plot(Y1,Y2)
mp.title("Representation de y2(t) en fonction de y1(t)")
mp.xlabel("y1(t)")
mp.ylabel("y2(t)")
mp.savefig("second-equation-function.png")




#------ Reprentation du champ des tangentes --------#

f = function_exp_arctan()

plot_tang_field(-4.,4.,-4.,4.,20,20,f[0])

#----- Comparaison entre les differentes methodes -----#

def test_convergence(y0,t0,tf,n,f,g,meth):
    
    
    X = np.arange(2,n+1,2)
    Y = np.zeros([X.shape[0],1])
    Y.shape =  (X.shape[0])
    k = 0
    for i in X:
        y = y0
        t = t0
        diff_max = 0.
        for j in range(i):
 
            y = meth(y,t,(tf-t0)/i,f)
       
            S = g(t)
       
            if( lg.norm(y-S) > diff_max):
                diff_max = lg.norm(y-S)
                
            t = t+(tf-t0)/i
 
        Y[k] = diff_max
        k = k+1

    return (X,Y)


def solution(t):
    return [math.exp(math.atan(t))]


f = function_exp_arctan()


(X,Y1) = test_convergence([1.],0.,20.,100,f,solution,step_euler)
Y2 = test_convergence([1.],0.,20.,100,f,solution,step_point_milieu)[1]
Y3 = test_convergence([1.],0.,20.,100,f,solution,step_heun)[1]
Y4 = test_convergence([1.],0.,20.,100,f,solution,runge_kutta_step)[1]

mp.clf()

mp.title("Convergence des differentes methodes")
mp.xlabel("Indice du pas")
mp.ylabel("Difference absolue de la solution exacte")

a =mp.plot(X,Y1,'b')
b =mp.plot(X,Y2,'r')
c =mp.plot(X,Y3,'g')
d =mp.plot(X,Y4,'k')
mp.legend((a,b,c,d),("Euler","Point milieu","Heun","Runge-Kutta"))
mp.savefig("convergence_comparaison.png")


